import os
import cv2 as cv
import pandas as pds
import numpy as np
import time
import random
from matplotlib import pyplot as plt
import json
import PIL
from PIL import Image
import shutil
from matplotlib import pyplot as plt
import cv2
import json

"""
清理保留文件，并按照规则文件将拆解
"""

#按照指定图像大小调整尺寸
def resize_image(image, height, width):
     top, bottom, left, right = (0, 0, 0, 0)

     #获取图像尺寸
     h, w, _ = image.shape
     #对于长宽不相等的图片，找到最长的一边
     longest_edge = max(h, w)    

     #计算短边需要增加多上像素宽度使其与长边等长
     if h < longest_edge:
         dh = longest_edge - h
         top = dh // 2
         bottom = dh - top
     elif w < longest_edge:
         dw = longest_edge - w
         left = dw // 2
         right = dw - left
     else:
         pass 

     #RGB颜色
     BLACK = [0, 0, 0]
     #给图像增加边界，是图片长、宽等长，cv2.BORDER_CONSTANT指定边界颜色由value指定
     constant = cv2.copyMakeBorder(image, top , bottom, left, right, cv2.BORDER_CONSTANT, value = BLACK)

     #调整图像大小并返回
     return cv2.resize(constant, (height, width))



def readTxt(path):
    with open(path,"r",encoding="utf-8") as fp:
        return fp.read()


rootdir="/home/gis/gisdata/data/jupyterlabhub/gitcode/cardanger/dataset/newcardata" # 根目录

imgdir=os.path.join(rootdir,"images")
removedir=os.path.join(rootdir,"removeimg")
splitdir=os.path.join(rootdir,"splitimages")
boxdir=os.path.join(splitdir,"boxs") # 框选图片
labelsdir=os.path.join(splitdir,"labels") # 对应的文件
csvpath=os.path.join(splitdir,"img_label.csv") 
labellist=os.listdir(labelsdir)  # 标注列表
imglist=os.listdir(imgdir) # 图片列表
csvobj=pds.read_csv(csvpath)
# 建立一个保存文件的位置
clipimgpath=os.path.join(splitdir,"clipimgs")
if os.path.exists(clipimgpath):
    shutil.rmtree(clipimgpath)
os.mkdir(clipimgpath)



result={"id":[],"annotation":[],"images":[],"sig":[],"class":[]}
reflect={"source":[],"target":[],"label":[]}
# 裁切图像，并生成对应的csv文件
tempresult=[]
n=len(csvobj)
imgcount=0
for i in range(n):
    dager=int(csvobj.loc[i,"danger"])
    if dager<0 :
        continue
    labelname=csvobj.loc[i,"label"]
    pcname=csvobj.iloc[i,2]
    pcname=pcname.split('\\')[-1]
    if not os.path.exists(os.path.join(boxdir,pcname)): 
        continue
    label= readTxt( os.path.join(labelsdir,labelname)) # 获取文件
    newclipimgname=os.path.basename(os.path.join(labelsdir,labelname)).split(".")[0]
    label=json.loads(label) # 解析文件
    imgname,[xl,yl,xr,yr],label,confidence=label["imgname"],label["box"],label["label"],label["confidence"]
    if label!="car":
        continue
    # 重新调整图片大小
    w,h=xr-xl,yr-yl
    dx,dy=0.0*w,0.0*h
    xl,yl,xr,yr=int(xl+dx),int(yl+dy),int(xr-dx),int(yr-dy) # 重新调整

    img=np.array(Image.open(os.path.join(imgdir,imgname)))
    h,w,c=img.shape
    xl=0 if xl<0 else xl
    yl=0 if yl<0 else yl
    xr=w-1 if xr>=w else xr
    yr=h-1 if yr>=h else yr
    w1,h1=xr-xl,yr-yl
    #if h1/w1>2 or w1/h1>4 or w1<=0 or h1 <=0  :
        #continue
    # 裁切图片
    chdimg=img[yl:yr,xl:xr,:] # 框
    h,w,c=chdimg.shape
    if h<64 and w<64:
        continue
    chdimg=resize_image(chdimg,256,256)
    # 保存裁切之后的图片
    chdimg=Image.fromarray(chdimg) # 图片保存
    #chdimg=chdimg.resize((256,256),resample=PIL.Image.LANCZOS)
    print(imgcount,h,w,chdimg.size)
    imgcount=imgcount+1
    chdimg.save(os.path.join(clipimgpath,"{}.jpg".format(newclipimgname))) #定义保存的文件格式
    try:
        testtempimg=Image.open(os.path.join(clipimgpath,"{}.jpg".format(newclipimgname)))
    except:
        os.remove(os.path.join(clipimgpath,"{}.jpg".format(newclipimgname)))
        continue #数据不能正常读取,删除对应的文件
    reflect["source"].append(pcname)
    reflect["target"].append(newclipimgname)
    reflect["label"].append(imgname)
    tempmeta={"id":csvobj.loc[i,"id"],
            "annotation":os.path.join(labelsdir,labelname),
            "images":os.path.join(clipimgpath,"{}.jpg".format(newclipimgname)),
            "class":str(csvobj.loc[i,"danger"])}
    tempresult.append(tempmeta)

reflectcsv=pds.DataFrame(reflect)
# 生成对应的文件格式
trsig=0
ttsig=0
vlsig=0
for temp in tempresult:
    result["id"].append(temp["id"])
    result["annotation"].append(temp["annotation"])
    result["images"].append(temp["images"])
    sig=random.random()
    if sig<1:
        result["sig"].append("train")
        trsig=trsig+1
    elif sig<0.95:
        result["sig"].append("test")
        ttsig=ttsig+1
    elif sig<=1:
        result["sig"].append("vail")
        vlsig=vlsig+1
    # 图片分类
    result["class"].append(temp["class"])

print(ttsig,vlsig,trsig)
result = pds.DataFrame(result)
result.to_csv(os.path.join(rootdir,"train_test_vail.csv"),encoding="utf-8")
reflectcsv.to_csv(os.path.join(rootdir,"reflectcsv.csv"),encoding="utf-8")
print(os.path.join(rootdir,"train_test_vail.csv"))
# 压缩文件
localtime = time.localtime(time.time())
shelltext="zip  -r  ./clipimgs{}_{}.zip  {} ".format(localtime.tm_mon,localtime.tm_mday,clipimgpath)
print('文件处理完成，正在进行文件压缩命令 :{} ，请勿关闭软件。。。。'.format(shelltext))
os.system(shelltext)
print(os.path.join(rootdir,"train_test_vail.csv"))
print(imgcount)
print(ttsig,vlsig,trsig)
'''
for labelname in labellist:
    label= readTxt( os.path.join(labelsdir,labelname)) # 获取文件
    label=json.loads(label) # 解析文件
    imgname,box,label,confidence=label["imgname"],label["box"],label["label"],label["confidence"]
    box=[box[0]-0.3*box[2],box[1]-0.3*box[3],box[2]+0.6*box[2],box[3]+0.3*box[3]]
    box=int(box)
    img=np.array(Image.open(os.path.join(imgdir,imgname)))
    h,w,c=img.shape
    box[0]=0 if box[0]<0 else box[0]
    box[1]=0 if box[1]<0 else box[1]
    box[2]=h-box[0] if box[2]+box[0]>=h else box[2]
    box[3]=w-box[1] if box[3]+box[1]>=w else box[3]
    # 裁切图片
    chdimg=img[box[0]:box[0]+box[2],box[1]:box[1]+box[3],:] # 框
    # 保存裁切之后的图片
    chdimg=Image.fromarray(chdimg) # 图片保存
    chdimg.save()
    pass
'''